##########################################################################################

library(ktplots)
#library(Seurat)
#library(SeuratDisk)
library(optparse)

##########################################################################################

option_list <- list(
    make_option(c("--pvals_file"), type = "character"),
    make_option(c("--means_file"), type = "character"),
    make_option(c("--out_path"), type = "character")
)

if(1!=1){

  work_dir <- "~/20220915_gastric_multiple/dna_combinePublic"

  pvals_file <- "~/20220915_gastric_multiple/dna_combinePublic/finalPlot/revise/im_favored/immuneCell_cpdb/statistical_analysis_pvalues_IM_MUC6.txt"
  means_file <- "~/20220915_gastric_multiple/dna_combinePublic/finalPlot/revise/im_favored/immuneCell_cpdb/statistical_analysis_means_IM_MUC6.txt"
  #single_cell_file <- "~/20220915_gastric_multiple/dna_combinePublic/images/singleCell_MUC6/cpdb/IM_MUC6.h5seurat"
  out_path <- "~/20220915_gastric_multiple/dna_combinePublic/finalPlot/revise/im_favored/immuneCell_cpdb/"

}

##########################################################################################

parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

out_path <- opt$out_path
means_file <- opt$means_file
pvals_file <- opt$pvals_file

dir.create(out_path , recursive = T)

##########################################################################################

#scRNA <- LoadH5Seurat(single_cell_file)
pvals <- read.delim(pvals_file, check.names = FALSE)
means <- read.delim(means_file, check.names = FALSE)

##########################################################################################
## 热图

cell_order <- c("Plasma" , "Pit_Mut" , "Pit_Other" , "Naive")

pvalue.threshold = 0.05
order.of.celltype = cell_order[length(cell_order):1]
ccc.number.max = 10
size.of.text = c(18,14,14,12)
color.palette = c("#4393C3","#ffdbba","#B2182B")

#library(tidyverse)
#library(RColorBrewer)
#library(scales)

pvalues=read.table(pvals_file,header = T,sep = "\t",stringsAsFactors = F,check.names = F)
pvalues=pvalues[,12:dim(pvalues)[2]] #此时不关注前11列
statdf=as.data.frame(colSums(pvalues < pvalue.threshold)) #统计在某一种细胞pair的情况之下，显著的受配体pair的数目；阈值可以自己选
colnames(statdf)=c("number")

#排在前面的分子定义为indexa；排在后面的分子定义为indexb
statdf$indexb=stringr::str_replace(rownames(statdf),"^.*\\|","")
statdf$indexa=stringr::str_replace(rownames(statdf),"\\|.*$","")
statdf$total_number=0

# 细胞相互间的交互加起来
for (i in 1:dim(statdf)[1]) {
  tmp_indexb=statdf[i,"indexb"]
  tmp_indexa=statdf[i,"indexa"]
  if (tmp_indexa == tmp_indexb) {
    statdf[i,"total_number"] = statdf[i,"number"]
  } else {
    statdf[i,"total_number"] = statdf[statdf$indexb==tmp_indexb & statdf$indexa==tmp_indexa,"number"]+
      statdf[statdf$indexa==tmp_indexb & statdf$indexb==tmp_indexa,"number"]
  }
}

#设置合适的细胞类型的顺序
if(is.null(order.of.celltype)){
	rankname=sort(unique(statdf$indexa))
} else {
	rankname=order.of.celltype
}

#转成因子类型，画图时，图形将按照预先设置的顺序排列
statdf$indexa=factor(statdf$indexa,levels = rankname[length(rankname):1])
statdf$indexb=factor(statdf$indexb,levels = rankname)
#statdf$indexa=factor(statdf$indexa,levels = rankname)

#调整图上数值的范围
if(is.null(ccc.number.max)){
	limits.max = ceiling(max(statdf$total_number) / 10) * 10
}else{
	statdf$total_number[statdf$total_number > ccc.number.max] = ccc.number.max
	limits.max = ccc.number.max
}

statdf <- subset( statdf , indexa != "Tumor" & indexb != "Tumor")

p <-ggplot(statdf,aes(x=indexa,y=indexb,fill=total_number))+
	geom_tile(color="white")+
	scale_fill_gradientn(colours = color.palette,limits=c(0,limits.max) , name="number of interactions")+
	scale_x_discrete("")+
	scale_y_discrete("")+
	theme_minimal()+
	theme(
	  axis.title = element_text(size = size.of.text[1]),
	  axis.text.x.bottom = element_text(hjust = 1, vjust = NULL, angle = 45,size = 14,color = "black",face='bold'),
	  axis.text.y.left = element_text(size = 14,color = "black",face='bold'),
	  legend.title = element_text(size = size.of.text[3]),
	  #legend.text = element_text(size = size.of.text[4]),
	  legend.position = "right" ,
	  panel.grid = element_blank()
	)

out_name <- paste0( out_path , "/IM_MUC6.immuneCell.heatmap.pdf" )
ggsave( out_name , p , width = 8 , height = 6 )

##########################################################################################

a_nocomplex <- grep( "complex" , pvals$partner_a , invert = T)
b_nocomplex <- grep( "complex" , pvals$partner_b , invert = T)

use_index <- a_nocomplex[a_nocomplex %in% b_nocomplex]

pvals_use <- pvals[use_index,]
means_use <- means[use_index,]

neg_log10_th= -log10(0.05) #前面画第一张图的时候用的什么值当作显著性阈值，这里保持一致
means_exp_log2_th=0 #表达均值也限制一下
notused.cell=NULL #确定不包含的细胞，默认是空值
used.cell=NULL #必须含有的细胞，默认是空值
neg_log10_th2=3 #限定显著性的最大值
means_exp_log2_th2=c(-2,0.75) #限定表达值的范围
cell.pair=NULL #这里是自定义的顺序，若是可选细胞对的子集，则只展示子集，若有交集则只展示交集；空值情况下，会根据可选细胞对自动排序
cell.pair=c(
  "Plasma|Plasma","Plasma|Pit_Mut","Plasma|Pit_Other",
	"Pit_Mut|Plasma","Pit_Mut|Pit_Mut" , "Pit_Mut|Pit_Other" ,
	"Pit_Other|Plasma" , "Pit_Other|Pit_Mut" , "Pit_Other|Pit_Other"
	)

cell.pair=c(
  "Plasma|Pit_Mut","Plasma|Pit_Other","Pit_Mut|Plasma","Pit_Other|Plasma" 
  )

gene.pair=NULL #作用同上
color_palette = c("#313695", "#4575B4", "#ABD9E9", "#FFFFB3", "#FDAE61", "#F46D43", "#D73027", "#A50026")
text_size = 12

#-----------------------------------------------------------
pvalues=pvals_use
pvalues=pvalues[,c(2,12:dim(pvalues)[2])]
RMpairs=names(sort(table(pvalues$interacting_pair))[sort(table(pvalues$interacting_pair)) > 1])
pvalues=pvalues[!(pvalues$interacting_pair %in% RMpairs),]
pvalues.df1=reshape2::melt(pvalues,id="interacting_pair")
colnames(pvalues.df1)=c("geneA_geneB","cellA_cellB","pvalue")
pvalues.df1$neg_log10=-log10(pvalues.df1$pvalue)
pvalues.df1$geneA_geneB_cellA_cellB=paste(pvalues.df1$geneA_geneB,pvalues.df1$cellA_cellB,sep = ",")

means=means_use
means=means[,c(2,12:dim(means)[2])]
rmpairs=names(sort(table(means$interacting_pair))[sort(table(means$interacting_pair)) > 1])
means=means[!(means$interacting_pair %in% rmpairs),]
means.df1=reshape2::melt(means,id="interacting_pair")
colnames(means.df1)=c("geneA_geneB","cellA_cellB","means_exp")
means.df1$geneA_geneB_cellA_cellB=paste(means.df1$geneA_geneB,means.df1$cellA_cellB,sep = ",")
means.df1=means.df1[,c("geneA_geneB_cellA_cellB","means_exp")]

raw.df=merge(pvalues.df1,means.df1,by="geneA_geneB_cellA_cellB")
raw.df$means_exp_log2=log2(raw.df$means_exp)
## 不做log转化
raw.df$means_exp_log2=raw.df$means_exp

#-----------------------------------------------------------
#根据第一组阈值筛选
raw.df <- subset( raw.df , cellA_cellB %in% cell.pair )
final.df=dplyr::filter(raw.df,neg_log10 > neg_log10_th & means_exp_log2 > means_exp_log2_th)

final.df$geneA=stringr::str_replace(final.df$geneA_geneB,"_.*$","") #此处的geneA geneB有不妥之处，
final.df$geneB=stringr::str_replace(final.df$geneA_geneB,"^.*_","") #没有考虑到complex的命名规则
final.df$cellA=stringr::str_replace(final.df$cellA_cellB,"\\|.*$","") #后面尽量不用
final.df$cellB=stringr::str_replace(final.df$cellA_cellB,"^.*\\|","")

  #-----------------------------------------------------------
  #根据其它规则过滤
  #有明确不呈现在图中的细胞，过滤如下
  if (!is.null(notused.cell)) {
    final.df=final.df[!(final.df$cellA %in% notused.cell),]
    final.df=final.df[!(final.df$cellB %in% notused.cell),]
  }
  #两种细胞相同的话一般不展示
  final.df=final.df[!(final.df$cellA==final.df$cellB),]

  #-----------------------------------------------------------
  #提取原始矩阵
  final.df.gene=unique(final.df$geneA_geneB)
  final.df.cell=unique(final.df$cellA_cellB)
  #pair中必须含有的细胞
  if (!is.null(used.cell)){
    tmp_cell=c()
    for (i in used.cell) {
      tmp_cell=union(tmp_cell,final.df.cell[str_detect(final.df.cell,i)])
    }
    final.df.cell=tmp_cell
  }
  raw.df=raw.df[raw.df$geneA_geneB %in% final.df.gene, ]
  raw.df=raw.df[raw.df$cellA_cellB %in% final.df.cell, ]

  #-----------------------------------------------------------
  #范围修正
  raw.df$neg_log10=ifelse(raw.df$neg_log10 > neg_log10_th2,neg_log10_th2,raw.df$neg_log10)
  raw.df$means_exp_log2=ifelse(
    raw.df$means_exp_log2 > means_exp_log2_th2[2],
    means_exp_log2_th2[2],
    ifelse(
      raw.df$means_exp_log2 < means_exp_log2_th2[1],
      means_exp_log2_th2[1],
      raw.df$means_exp_log2
    )
  )

  #-----------------------------------------------------------
  #cellA-cellB的排列顺序
  raw.df$cellA_cellB=as.character(raw.df$cellA_cellB)
  if (!is.null(cell.pair)) {
    tmp_pair=intersect(cell.pair,unique(raw.df$cellA_cellB))
    raw.df=raw.df[raw.df$cellA_cellB %in% tmp_pair,]
    raw.df$cellA_cellB=factor(raw.df$cellA_cellB,levels = tmp_pair)
  } else {
    tmp_pair=sort(unique(raw.df$cellA_cellB))
    raw.df$cellA_cellB=factor(raw.df$cellA_cellB,levels = tmp_pair)
  }
  #geneA-geneB的排列顺序
  raw.df$geneA_geneB=as.character(raw.df$geneA_geneB)
  if (!is.null(gene.pair)) {
    tmp_pair=intersect(gene.pair,unique(raw.df$geneA_geneB))
    raw.df=raw.df[raw.df$geneA_geneB %in% tmp_pair,]
    raw.df$geneA_geneB=factor(raw.df$geneA_geneB,levels = tmp_pair)
  } else {
    tmp_pair=sort(unique(raw.df$geneA_geneB))
    raw.df$geneA_geneB=factor(raw.df$geneA_geneB,levels = tmp_pair)
  }

p <- ggplot(raw.df,aes(cellA_cellB,geneA_geneB))+
    geom_point(aes(size=neg_log10,color=means_exp_log2))+
    scale_color_gradientn("scaled_means",colors = color_palette)+
    scale_size_continuous("-log10(p value)")+
    theme_bw()+
    theme(
      panel.grid.major=element_blank(),
      panel.grid.minor=element_blank(),
      panel.background = element_blank(),
      panel.border = element_blank(),
      axis.title = element_blank(),
      axis.text.x.bottom = element_text(hjust = 1, angle = 45, size=10, color = "black" , face = "bold"),
      axis.text.y.left = element_text(size = 10,color = "black" , face = "bold"),
      axis.line = element_line(size = 0.5),
      axis.ticks.length = unit(0.15,"cm") 
    )

out_name <- paste0( out_path , "/IM_MUC6.immuneCell.cpdb.pdf" )
ggsave( out_name , p , width = 4 , height = 4 )

